package edu.buaa.common.client;

import com.google.common.util.concurrent.ListenableFuture;
import edu.buaa.common.transaction.AbstractTransaction;
import edu.buaa.common.utils.PFieldList;
import edu.buaa.common.utils.PVal;
import edu.buaa.common.utils.SQLSchemaDB;
import edu.buaa.common.transaction.*;
import edu.buaa.utils.Pair;
import edu.buaa.utils.Triple;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.*;

public class PostgreSQLTimePointClient extends PostgreSqlExecutorClient {

    private final SQLSchemaDB schema;

    public PostgreSQLTimePointClient(String serverHost, String dbName, int parallelCnt, SQLSchemaDB schema) throws Exception {
        super(serverHost, dbName, parallelCnt);
        this.schema = schema;
    }

    @Override
    protected List<String> createTables() {
        return Arrays.asList(
                createTable("node", schema.nodeStatic),
                createTable("node_tp", schema.nodeTemporal),
                createTable("rel", schema.relStatic),
                createTable("rel_tp", schema.relTemporal)
        );
    }

    @Override
    public ListenableFuture<ServerResponse> execute(AbstractTransaction tx) throws InterruptedException {
        switch (tx.getTxType()) {
            case tx_import_static_data:
                return this.submit(execute((ImportStaticDataTx) tx));
            case tx_import_temporal_data:
                return this.submit(execute((ImportTemporalDataTx) tx), tx.getSection());
            case tx_query_snapshot:
                return this.submit(execute((SnapshotQueryTx) tx));
            case tx_update_temporal_data:
                return this.submit(execute((UpdateTemporalDataTx) tx));
            case tx_query_snapshot_aggr_max:
                return this.submit(execute((SnapshotAggrMaxTx) tx));
            case tx_query_snapshot_aggr_duration:
                return this.submit(execute((SnapshotAggrDurationTx) tx));
            case tx_query_entity_history:
                return this.submit(execute((EntityHistoryTx) tx));
            case tx_query_road_by_temporal_condition:
                return this.submit(execute((EntityTemporalConditionTx) tx));
            default:
                throw new UnsupportedOperationException();
        }
    }

    private void setProp(PreparedStatement stat, int pos, PVal val) throws SQLException {
        switch (val.getType()){
            case STRING:
                stat.setString(pos, val.s());
                break;
            case INT:
                stat.setInt(pos, val.i());
                break;
            case FLOAT:
                stat.setFloat(pos, val.f());
                break;
        }
    }

    private PVal setPVal(ResultSet rs, ResultSetMetaData rsm, int pos) throws SQLException {
        String type = rsm.getColumnTypeName(pos).toUpperCase();
        if (type.startsWith("FLOAT")){
            return PVal.f(rs.getFloat(pos));
        }else if(type.startsWith("INT")){
            return PVal.i(rs.getInt(pos));
        }else if(type.contains("CHAR")){
            return PVal.s(rs.getString(pos));
        }
        throw new IllegalStateException("got type "+type+" on column.");
    }

    private String insertSQL(String tableName, List<String> props){
        return "INSERT INTO "+tableName+"("+String.join(",", props)+") VALUES ("+qMarks(props.size())+")";
    }

    private String qMarks(int size){
        String[] arr = new String[size];
        Arrays.fill(arr, "?");
        return String.join(",", arr);
    }

    private Req execute(ImportStaticDataTx tx) {
        return conn -> {
            conn.setAutoCommit(false);
            PFieldList nodesData = tx.getNodes();
            List<String> nProps = new ArrayList<>(nodesData.keys());
            int nSize = nodesData.size();
            PFieldList relData = tx.getRels();
            int rSize = relData.size();
            List<String> rProps = new ArrayList<>(relData.keys());
            try (PreparedStatement stat1 = conn.prepareStatement(insertSQL("node", nProps));
                 PreparedStatement stat2 = conn.prepareStatement(insertSQL("rel", rProps))) {
                for (int i=0; i<nSize; i++) {
                    for(int j=0; j<nProps.size(); j++){
                        String key = nProps.get(j);
                        setProp(stat1, j, nodesData.get(key, i));
                    }
                    stat1.addBatch();
                }
                stat1.executeBatch();
                for (int i=0; i<rSize; i++) {
                    for(int j=0; j<rProps.size(); j++){
                        String key = rProps.get(j);
                        setProp(stat2, j, relData.get(key, i));
                    }
                    stat2.addBatch();
                }
                stat2.executeBatch();
                conn.commit();
            } catch (SQLException e) {
                throw new TransactionFailedException(e, tx);
            } finally {
                conn.setAutoCommit(true);
            }
            return new AbstractTransaction.Result();
        };
    }

    private Req execute(ImportTemporalDataTx tx) {
        return conn -> {
            String tableName = tx.isNode()?"node_tp":"rel_tp";
            PFieldList data = tx.getData();
            List<String> props = new ArrayList<>(data.keys());
            int tSize = data.size();
            try (PreparedStatement stmt = conn.prepareStatement(insertSQL(tableName, props))) {
                conn.setAutoCommit(false);
                for (int i=0; i<tSize; i++) {
                    for(int j=0; j<props.size(); j++){
                        String key = props.get(j);
                        setProp(stmt, j, data.get(key, i));
                    }
                    stmt.addBatch();
                }
                stmt.executeBatch();
                conn.commit();
            } catch (SQLException e) {
                throw new TransactionFailedException(e, tx);
            } finally {
                conn.setAutoCommit(true);
            }
            return new AbstractTransaction.Result();
        };
    }

    private Req execute(UpdateTemporalDataTx tx) {
        return conn -> {
            conn.setAutoCommit(false);
            String tableName = tx.isNode()?"node_tp":"rel_tp";
            PFieldList data = tx.getData();
            List<String> props = new ArrayList<>(data.keysWithout("entity","start","end"));
            int tSize = data.size();
            String propStr = String.join(",", props);

            try (PreparedStatement stmt1 = conn.prepareStatement("select entity, max(t) as max_t from "+tableName+" where entity=? and t<=? group by entity");
                 PreparedStatement stmt2 = conn.prepareStatement("select "+propStr+" from "+tableName+" where entity=? and t=?");
                 PreparedStatement stmt3 = conn.prepareStatement("insert into "+tableName+"(t, entity, "+propStr+") values("+qMarks(props.size()+2)+") " + "on CONFLICT (entity, t) DO NOTHING");
                 PreparedStatement stmt4 = conn.prepareStatement("delete from "+tableName+" where entity=? and t>=? and t<=?")) {
                for (int k=0; k<tSize; k++) {
                    String entityId = data.get("entity", k).s();
                    int beginT = data.get("start", k).i();
                    int endT = data.get("end", k).i();

                    stmt1.setString(1, entityId);
                    stmt1.setInt(2, endT);
                    ResultSet rs = stmt1.executeQuery();
                    // insert into table from t1 + 1 with the old value.
                    if (rs.next()) {
                        stmt2.setString(1, rs.getString("entity"));
                        stmt2.setInt(2, rs.getInt("max_t"));
                        ResultSet r = stmt2.executeQuery();
                        if (r.next()) {
                            stmt3.setInt(1, endT + 1);
                            stmt3.setString(2, entityId);
                            int j = 3;
                            for (String prop : props) {
                                stmt3.setObject(j, r.getObject(prop));
                                j++;
                            }
                            stmt3.execute();
                        }
                    }
                    // delete the [t0, t1] rows.
                    stmt4.setString(1, entityId);
                    stmt4.setInt(2, beginT);
                    stmt4.setInt(3, endT);
                    stmt4.execute();
                    // insert into table from t0 with the new value.
                    stmt3.setInt(1, beginT);
                    stmt3.setString(2, entityId);
                    for (int i = 0; i < props.size(); i++) {
                        setProp(stmt3, i+3, data.get(props.get(i), k));
                    }
                    stmt3.execute();
                    conn.commit();
                }
            } catch (SQLException | NullPointerException e) {
                throw new TransactionFailedException(e, tx);
            }
            return new AbstractTransaction.Result();
        };
    }

    private Req execute(SnapshotQueryTx tx) {
        return conn -> {
            String tableName = tx.isNode()?"node_tp":"rel_tp";
            List<Pair<String, PVal>> res = new ArrayList<>();
            try (PreparedStatement stmt1 = conn.prepareStatement("select ts.entity, " + tx.getPropertyName() + " from "+tableName+" as ts, " +
                         "(select entity, max(t) as max_t from "+tableName+" where t <= ? group by entity) as tmp where ts.entity = tmp.entity and ts.t = tmp.max_t")) {
                conn.setAutoCommit(false);
                int snapshotTime = tx.getTimestamp();
                stmt1.setInt(1, snapshotTime);
                ResultSet rs = stmt1.executeQuery();
                ResultSetMetaData rsm = rs.getMetaData();
                while (rs.next()) {
                    res.add(Pair.of(rs.getString("entity"), setPVal(rs, rsm, 2)));
                }
            } catch (SQLException e) {
                throw new TransactionFailedException(e, tx);
            } finally {
                conn.setAutoCommit(true);
            }
            SnapshotQueryTx.Result result = new SnapshotQueryTx.Result();
            result.setStatus(res);
            return result;
        };
    }

    private Req execute(SnapshotAggrMaxTx tx) {
        return conn -> {
            String tableName = tx.isNode()?"node_tp":"rel_tp";
            List<Pair<String, PVal>> res = new ArrayList<>();
            try (PreparedStatement stmt2 = conn.prepareStatement("select ts.entity, max(" + tx.getP() + ") as max_p from "+tableName+" as ts, " +
                    "(select entity, max(t) as max_t from "+tableName+" where t<=? group by entity) as tmp " +
                    "where tmp.entity=ts.entity and tmp.max_t<=ts.t and ts.t<=? group by ts.entity")) {
                conn.setAutoCommit(false);
                stmt2.setInt(1, tx.getT0());
                stmt2.setInt(2, tx.getT1());
                ResultSet rs = stmt2.executeQuery();
                HashMap<String, PVal> st = new HashMap<>();
                ResultSetMetaData rsm = rs.getMetaData();
                while (rs.next()) {
                    res.add(Pair.of(rs.getString("entity"), setPVal(rs, rsm, 2)));
                }
            } catch (SQLException e) {
                throw new TransactionFailedException(e, tx);
            } finally {
                conn.setAutoCommit(true);
            }
            SnapshotAggrMaxTx.Result result = new SnapshotAggrMaxTx.Result();
            result.setPropMaxValue(res);
            return result;
        };
    }

    private Req execute(SnapshotAggrDurationTx tx) {
        final TreeSet<PVal> intStart = tx.getIntStartTreeSet();
        return conn -> {
            List<Triple<String, PVal, Integer>> res = new ArrayList<>();
            try (PreparedStatement stmt0 = conn.prepareStatement("select node_id, max(t) as real_st from temporal_status where t <= ? group by node_id");
                 PreparedStatement stmt1 = conn.prepareStatement("select t, " + tx.getP() + " as property from temporal_status where node_id = ? and t >= ? and t <= ? order by t")) {
                conn.setAutoCommit(false);
                stmt0.setInt(1, tx.getT0());
                ResultSet rs = stmt0.executeQuery();
                while (rs.next()) {
                    int id = rs.getInt("node_id");
                    stmt1.setInt(1, id);
                    stmt1.setInt(2, rs.getInt("real_st"));
                    stmt1.setInt(3, tx.getT1());
                    ResultSet r = stmt1.executeQuery();
                    int lastTime = -1;
                    float lastStatus = Float.MIN_VALUE;
                    HashMap<Float, Integer> buffer = new HashMap<>();
                    while (r.next()) {
                        float property = r.getFloat("property");
                        int t = r.getInt("t");
                        if (lastTime != -1) {
                            int duration = t - lastTime;
                            Float statusGrp = intStart.floor( lastStatus);
                            if(statusGrp!=null) buffer.merge(statusGrp, duration, Integer::sum);
                            lastTime = t;
                        } else {
                            lastTime = tx.getT0();
                        }
                        lastStatus = property;
                    }
                    if (lastTime != -1) {
                        Float statusGrp = intStart.floor(lastStatus);
                        if(statusGrp!=null) buffer.merge(statusGrp, tx.getT1() - lastTime + 1, Integer::sum);
                    }
                    for (Map.Entry<Float, Integer> e : buffer.entrySet()) {
                        res.add(Triple.of(id, e.getKey(), e.getValue()));
                    }
//                        System.out.println(res.size());
                }
            } catch (SQLException e) {
                throw new TransactionFailedException(e, tx);
            } finally {
                conn.setAutoCommit(true);
            }
            SnapshotAggrDurationTx.Result result = new SnapshotAggrDurationTx.Result();
            result.setStatusDuration(res);
            return result;
        };
    }

    private Req execute(EntityHistoryTx tx) {
        return conn -> {
            List<Triple<Integer, Integer, Float>> res = new ArrayList<>();
            try (PreparedStatement stmt0 = conn.prepareStatement("select max(t) as real_st from temporal_status where t <= ? and node_id=?");
                 PreparedStatement stmt1 = conn.prepareStatement("select t, " + tx.getProp() + " as property from temporal_status where node_id = ? and t >= ? and t <= ? order by t")) {
                conn.setAutoCommit(false);
                stmt0.setInt(1, tx.getBeginTime());
                stmt0.setInt(2, tx.getEntity());
                ResultSet rs = stmt0.executeQuery();
                if (rs.next()) {
                    stmt1.setInt(1, tx.getEntity());
                    stmt1.setInt(2, rs.getInt("real_st"));
                    stmt1.setInt(3, tx.getEndTime());
                    ResultSet r = stmt1.executeQuery();
                    int lastTime = -1;
                    float lastStatus = Float.MIN_VALUE;
                    while (r.next()) {
                        float property = r.getFloat("property");
                        int t = r.getInt("t");
                        if (lastTime != -1) {
                            res.add(Triple.of(lastTime, t-1, lastStatus));
                            lastTime = t;
                        } else {
                            lastTime = tx.getBeginTime();
                        }
                        lastStatus = property;
                    }
                    if (lastTime != -1) {
                        res.add(Triple.of(lastTime, tx.getEndTime(), lastStatus));
                    }
                }
            } catch (SQLException e) {
                throw new TransactionFailedException(e, tx);
            } finally {
                conn.setAutoCommit(true);
            }
            EntityHistoryTx.Result result = new EntityHistoryTx.Result();
            result.setHistory(res);
            return result;
        };
    }

    private Req execute(EntityTemporalConditionTx tx) {
        return conn -> {
            List<Integer> res = new ArrayList<>();
            try (PreparedStatement stmt0 = conn.prepareStatement("select distinct ts.node_id from temporal_status as ts, " +
                    "(select node_id, max(t) as max_t from temporal_status where t<=? group by node_id) as tmp " +
                    "WHERE tmp.node_id=ts.node_id AND ts."  + tx.getP() + " >= ? AND ts." + tx.getP() + " <= ? AND "+
                    "tmp.max_t<=ts.t AND ts.t<=?")) {
                conn.setAutoCommit(false);
                stmt0.setInt(1, tx.getT0());
                stmt0.setFloat(2, tx.getVMin());
                stmt0.setFloat(3, tx.getVMax());
                stmt0.setInt(4, tx.getT1());
                ResultSet rs = stmt0.executeQuery();
                while (rs.next()) {
                    res.add(rs.getInt("node_id"));
                }
            } catch (SQLException e) {
                throw new TransactionFailedException(e, tx);
            } finally {
                conn.setAutoCommit(true);
            }
            EntityTemporalConditionTx.Result result = new EntityTemporalConditionTx.Result();
            result.setEntities(res);
            return result;
        };
    }
}
